_{Comment @coderabbitai help to get the list of available commands and usage tips.}

/bot run

@CodeRabbit please review the piecewise related code.

#!/bin/bash
cat tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py

#!/bin/bash
cat tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py

@nvchenghaoz Here is a detailed review of the piecewise CUDA graph code.

Review: Piecewise CUDA Graph (torch_cudagraph.py + compile_model.py)

🔴 Issues

1. num_tokens not forwarded in PiecewiseCapturedGraph.forward fallback

num_tokens is explicitly extracted from **kwargs in the signature, so it is not present in **kwargs when original_model is called in the fallback branch.

def forward(self, *args, num_tokens: Optional[int] = None, **kwargs) -> Any:
    if self.split_gm is not None:
        ...
    return self.original_model(*args, **kwargs)  # num_tokens is silently dropped

If original_model (or any caller down the chain) relies on receiving num_tokens, this will silently break. The fix is straightforward:

return self.original_model(*args, num_tokens=num_tokens, **kwargs)

2. Silent exception swallow in _reconstruct_output

def _reconstruct_output(self, result: Any) -> Any:
    if not isinstance(result, tuple) or self._out_spec is None:
        return result
    try:
        return self._out_spec.unflatten(list(result))
    except Exception:
        return result  # ← exception silently discarded

Any TreeSpec/shape mismatch (e.g. a mismatched out_spec from the monolithic path) will be silently swallowed, returning a raw tuple instead of a structured ModelOutput. This is very hard to debug. At minimum, log a warning:

except Exception as e:
    ad_logger.warning("_reconstruct_output: failed to unflatten output (%s), returning raw tuple", e)
    return result

3. Address-unstable kwargs with no detected dynamic dim are silently skipped

In _allocate_static_input_buffers:

if dyn_dim is not None:
    self._static_input_buffers[key] = (torch.empty_like(v1), dyn_dim)

If a kwarg tensor has an unstable address across calls but its shape is static (same shape at max_bucket and max_bucket - 1), dyn_dim remains None and no static buffer is allocated. CUDA graph replay will then use the original (unstable) pointer and can crash or silently corrupt state. This should be either:

• Explicitly handled with torch.empty_like(v1) buffer and a direct full copy (no narrow), or
• Logged as a warning so it is visible:

else:
    ad_logger.warning(
        "PiecewiseCapturedGraph: kwarg '%s' has unstable address but no dynamic dim found — "
        "not buffered; CUDA graph replay may crash if address changes",
        key,
    )

🟡 Warnings

4. _truncate_output dim-matching heuristic is fragile for square tensors

for d in range(v.ndim):
    if v.shape[d] == bucket:
        return v.narrow(d, 0, num_tokens)

For a tensor with shape [512, 512] and bucket=512, this will truncate on dim=0 (the first match), which may be wrong if the token dimension is dim=1. Consider passing the known token dimension explicitly, or at least adding an assertion that only one dim matches:

matching = [d for d in range(v.ndim) if v.shape[d] == bucket]
if len(matching) > 1:
    ad_logger.warning("_truncate_output: ambiguous token dim for shape %s, bucket=%d", v.shape, bucket)

5. ADPiecewiseRunner.set_current_num_tokens(None) not in try/finally

In DualModeCapturedGraph.forward:

result = self.piecewise(*args, num_tokens=bucket, **kwargs)
ADPiecewiseRunner.set_current_num_tokens(None)   # ← not reached if piecewise raises
if bucket > num_tokens:
    result = self._truncate_output(result, num_tokens, bucket)
return result

If self.piecewise(...) raises, the global num_tokens state is left dirty, which could corrupt subsequent decode calls. Use try/finally:

try:
    result = self.piecewise(*args, num_tokens=bucket, **kwargs)
finally:
    ADPiecewiseRunner.set_current_num_tokens(None)

6. __getattr__ proxy in DualModeCapturedGraph can mask AttributeError on self.monolithic

def __getattr__(self, name: str):
    try:
        return super().__getattr__(name)
    except AttributeError:
        return getattr(self.monolithic.model, name)

If self.monolithic exists but self.monolithic.model raises AttributeError for name, the exception propagates without a clear message. Consider adding a fallback error message:

except AttributeError:
    pass
raise AttributeError(f"'{type(self).__name__}' and its inner model have no attribute '{name}'")

✅ Things that look good

• compile_model.py GraphModule traversal: The seen set + startswith guard correctly avoids double-compiling nested GMs. Setting full_model=mod for nested targets and None for the root is correct.
• _capture_inner_kwargs using a pre-hook: The pre-hook fires before the inner module's forward, so tensor references are fresh. The try/finally around handle.remove() correctly avoids hook leaks.
• CapturedGraph auto-detection of dynamic dims: Using probe_bs = max(1, batch_sizes[0] - 1) sensibly avoids OOB issues. The edge case of batch_sizes[0] == 1 (probe == max) safely defaults to dim=0.
• out_spec threading: Passing monolithic._out_spec into PiecewiseCapturedGraph after monolithic capture is a clean way to share the tree structure without re-running a warmup.
• Static buffer detection via data_ptr(): Using data_ptr() comparison across two calls is a robust way to detect address instability, better than static analysis of the graph.

@CodeRabbit more review please

#!/bin/bash
cat tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py
echo "====SEPARATOR===="
cat tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
echo "====SEPARATOR===="
cat tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py
echo "====SEPARATOR===="
cat tensorrt_llm/_torch/auto_deploy/utils/_graph.py
echo "====SEPARATOR===="
cat tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py

# In PiecewiseCapturedGraph.forward(...)
-    return self.original_model(*args, **kwargs)
+    return self.original_model(*args, num_tokens=num_tokens, **kwargs)

# In PiecewiseCapturedGraph._reconstruct_output(...)
     try:
         return self._out_spec.unflatten(list(result))
-    except Exception:
-        return result
+    except Exception as e:
+        ad_logger.warning("_reconstruct_output: failed to unflatten output (%s); returning raw tuple", e)
+        return result

# In PiecewiseCapturedGraph._allocate_static_input_buffers(...)
-    if dyn_dim is not None:
-        self._static_input_buffers[key] = (torch.empty_like(v1), dyn_dim)
+    if dyn_dim is not None:
+        self._static_input_buffers[key] = (torch.empty_like(v1), dyn_dim)
+    else:
+        # Address is unstable but no dynamic dim — still buffer and copy full tensor
+        self._static_input_buffers[key] = (torch.empty_like(v1), None)

# In PiecewiseCapturedGraph.forward(...) where copying kwargs into static buffers
-    buf, dyn_dim = self._static_input_buffers[key]
-    if dyn_dim is None:
-        continue
-    bs = args_batched[0].shape[dyn_dim]
-    buf.narrow(dyn_dim, 0, bs).copy_(v_current.narrow(dyn_dim, 0, bs))
-    kwargs_batched[key] = buf.narrow(dyn_dim, 0, bs)
+    buf, dyn_dim = self._static_input_buffers[key]
+    if dyn_dim is None:
+        buf.copy_(v_current)
+        kwargs_batched[key] = buf
+    else:
+        bs = args_batched[0].shape[dyn_dim]
+        buf.narrow(dyn_dim, 0, bs).copy_(v_current.narrow(dyn_dim, 0, bs))
+        kwargs_batched[key] = buf.narrow(dyn_dim, 0, bs)

# In DualModeCapturedGraph.forward(...)
-    result = self.piecewise(*args, num_tokens=bucket, **kwargs)
-    ADPiecewiseRunner.set_current_num_tokens(None)
+    try:
+        result = self.piecewise(*args, num_tokens=bucket, **kwargs)
+    finally:
+        ADPiecewiseRunner.set_current_num_tokens(None)

# In DualModeCapturedGraph.__getattr__(...)
-    try:
-        return super().__getattr__(name)
-    except AttributeError:
-        return getattr(self.monolithic.model, name)
+    try:
+        return super().__getattr__(name)
+    except AttributeError:
+        try:
+            return getattr(self.monolithic.model, name)
+        except AttributeError:
+            raise AttributeError(f\"{type(self).__name__} and its inner model have no attribute '{name}'\")

# In DualModeCapturedGraph._truncate_output(...)
-    for d in range(v.ndim):
-        if v.shape[d] == bucket:
-            return v.narrow(d, 0, num_tokens)
-    return v
+    matching = [d for d in range(v.ndim) if v.shape[d] == bucket]
+    if len(matching) != 1:
+        ad_logger.warning(\"_truncate_output: ambiguous dim for shape %s, bucket=%d; dims=%s\",
+                          tuple(v.shape), bucket, matching)
+    d = matching[0] if matching else -1
+    return v.narrow(d, 0, num_tokens) if d >= 0 else v

# In get_lm_head_node(...)
- if node.target is torch.ops.auto_deploy.trtllm_dist_all_gather:
+ if node.target in {
+     torch.ops.auto_deploy.trtllm_dist_all_gather.default,
+     torch.ops.auto_deploy.torch_dist_all_gather.default,
+ }:
      node = list(node.all_input_nodes)[0]

#!/bin/bash
set -euo pipefail

echo ">>> grep torch_cudagraph core spots"
rg -n --no-heading -C3 $'class PiecewiseCapturedGraph|def _reconstruct_output\\(|def _truncate_output\\(|class DualModeCapturedGraph|def __getattr__\\(|class CapturedGraph|dynamic_dims|set_current_num_tokens' tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py || true

echo ">>> show get_lm_head_node updates"
rg -n --no-heading -C4 'def get_lm_head_node|trtllm_dist_all_gather|torch_dist_all_gather' tensorrt_llm/_torch/auto_deploy/utils/_graph.py || true

@nvchenghaoz Thanks — I dug deeper into the piecewise path and related integration points. Findings and concrete patches below.

What looks good

• DualModeCapturedGraph now always clears ADPiecewiseRunner num_tokens via try/finally (Line 792-794).
• CapturedGraph auto-detects per-input dynamic dims and uses narrow() at replay; this is solid.
• The compile traversal that compiles nested GraphModules and passes full_model for inner-kwargs capture is sound.

Actionable fixes and improvements

1. Bug: num_tokens is dropped on eager fallback in PiecewiseCapturedGraph.forward

• Evidence: Line 650-651 returns original_model without forwarding num_tokens.
• Fix: pass it through.

Patch:

# tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py

class PiecewiseCapturedGraph(nn.Module):
    def forward(self, *args, num_tokens: Optional[int] = None, **kwargs) -> Any:
        if self.split_gm is not None:
            self._copy_to_static_buffers(kwargs)
            ADPiecewiseRunner.set_current_num_tokens(num_tokens)
            try:
                result = self.split_gm(*args, **kwargs)
            finally:
                ADPiecewiseRunner.set_current_num_tokens(None)
            return self._reconstruct_output(result)
-       return self.original_model(*args, **kwargs)
+       return self.original_model(*args, num_tokens=num_tokens, **kwargs)

2. Robustness: handle static-shaped but address-unstable kwargs in static buffer path

• Today static buffers are allocated only when a dynamic dim is detected. If a kwarg tensor’s address changes across calls but its shape is static, replay can still be unsafe without buffering.
• Change: allocate a buffer even when dyn_dim is None and do a full copy; use narrow only when a dyn_dim exists.

Patch (conceptual; adjust to your _copy_to_static_buffers/_allocate_static_input_buffers structure):

# When detecting instability:
if dyn_dim is not None:
    self._static_input_buffers[key] = (torch.empty_like(v1), dyn_dim)
else:
    # address unstable but static shape — still buffer it
    self._static_input_buffers[key] = (torch.empty_like(v1), None)

# When copying into buffers before replay:
buf, dyn_dim = self._static_input_buffers[key]
if dyn_dim is None:
    buf.copy_(v_current)
    kwargs_batched[key] = buf
else:
    bs = args_batched[0].shape[dyn_dim]
    buf.narrow(dyn_dim, 0, bs).copy_(v_current.narrow(dyn_dim, 0, bs))
    kwargs_batched[key] = buf.narrow(dyn_dim, 0, bs)

3. Safer truncation when bucket > real num_tokens

• Current _truncate_output finds the first dim equal to bucket (Line 740+ docstring). This is ambiguous for square-ish shapes (e.g., [B, S, S] or [S, S, H]).
• Minimal improvement: warn when multiple dims match and fall back gracefully.

Patch:

def _truncate_output(self, result: Any, num_tokens: int, bucket: int) -> Any:
    def cut(v: torch.Tensor) -> torch.Tensor:
        matching = [d for d in range(v.ndim) if v.shape[d] == bucket]
        if len(matching) != 1:
            ad_logger.warning(
                "_truncate_output: ambiguous dim for shape %s, bucket=%d; dims=%s",
                tuple(v.shape), bucket, matching
            )
        d = matching[0] if matching else -1
        return v.narrow(d, 0, num_tokens) if d >= 0 else v
    if isinstance(result, tuple):
        return tuple(cut(v) if isinstance(v, torch.Tensor) else v for v in result)
    return cut(result) if isinstance(result, torch.Tensor) else result

Optional but recommended: during capture, compare each output tensor’s shape at bucket vs probe to record its own token-dim index (e.g., self._output_dynamic_dims[i]), then use that instead of the heuristic.

4. Observability: log on output unflatten failure

• If _reconstruct_output uses a try/except around unflatten, add a warning to surface mismatches. If it doesn’t currently catch, consider adding this guard.

Patch:

def _reconstruct_output(self, result: Any) -> Any:
    if not isinstance(result, tuple) or self._out_spec is None:
        return result
    try:
        return self._out_spec.unflatten(list(result))
    except Exception as e:
        ad_logger.warning("_reconstruct_output: failed to unflatten output (%s); returning raw tuple", e)
        return result

5. Interop: unwrap both TRT-LLM and Torch all_gather around lm_head

• Evidence: get_lm_head_node currently unwraps trtllm_dist_all_gather only (Lines 581-583). For torch backend we should also handle torch_dist_all_gather.

Patch:

# tensorrt_llm/_torch/auto_deploy/utils/_graph.py

from torch.fx import GraphModule, Node
import torch

def get_lm_head_node(gm: GraphModule, output_node: Optional[Node] = None) -> Node:
    ...
-   if is_op(lm_head_node, torch.ops.auto_deploy.trtllm_dist_all_gather):
-       lm_head_node = lm_head_node.all_input_nodes[0]
+   if is_op(lm_head_node, {
+       torch.ops.auto_deploy.trtllm_dist_all_gather,
+       torch.ops.auto_deploy.torch_dist_all_gather,
+   }):
+       lm_head_node = lm_head_node.all_input_nodes[0]
    return lm_head_node

Nice-to-have

• DualModeCapturedGraph.getattr: if both wrapper and inner model lack the attribute, raise a clear message. If you want, I can supply a tiny tweak to avoid confusing AttributeError chaining.

Would you like me to stage these as a small follow-up patch and add a couple of focused unit tests (num_tokens passthrough and truncation-dim disambiguation)? If yes, I can prepare a plan next.

Close this one since it contains other work. This branch is just to check whether the change fixes the Qwen VLMs. Will create a new PR to have Piecewise CG changes only.